import numpy as np
import logging
from collections import deque
from typing import Optional, Tuple, Dict, List
from scipy import signal
from sklearn.decomposition import FastICA
import sys
import os

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if project_root not in sys.path:
    sys.path.append(project_root)

from response.BrainWaveResult import BrainWaveResult
from exception.base_exceptions import SignalProcessingException
from exception.exception_handler import global_exception_handler

logger = logging.getLogger(__name__)

class EOGArtifactRemover:
    """改进的EOG伪迹去除器，增强数值稳定性"""
    
    def __init__(self, 
                 blink_threshold: float = 50.0,  # 降低阈值
                 window_size: int = 50,
                 ica_components: int = 6,  # 适应8频域数据的ICA成分数量
                 sampling_rate: float = 512.0):
        """
        初始化改进的EOG伪迹去除器
        
        Args:
            blink_threshold: 眨眼检测阈值
            window_size: 滑动窗口大小
            ica_components: ICA成分数量
            sampling_rate: 采样率
        """
        self.blink_threshold = blink_threshold
        self.window_size = window_size
        self.ica_components = ica_components
        self.sampling_rate = sampling_rate
        
        # 初始化ICA模型，使用更稳定的参数
        self.ica_model = FastICA(
            n_components=ica_components,
            algorithm='parallel',
            whiten='unit-variance',  # 使用单位方差白化
            fun='logcosh',
            max_iter=200,  # 减少最大迭代次数
            tol=1e-3,  # 放宽收敛容忍度
            random_state=42  # 固定随机种子
        )
        
        # 信号缓冲区（8频域支持）
        self.signal_buffer = {
            'delta': deque(maxlen=window_size),
            'theta': deque(maxlen=window_size),
            # 8频域数据
            'low_alpha': deque(maxlen=window_size),
            'high_alpha': deque(maxlen=window_size),
            'low_beta': deque(maxlen=window_size),
            'high_beta': deque(maxlen=window_size),
            'low_gamma': deque(maxlen=window_size),
            'mid_gamma': deque(maxlen=window_size),
            # 向后兼容数据
            'alpha': deque(maxlen=window_size),
            'beta': deque(maxlen=window_size),
            'gamma': deque(maxlen=window_size),
            'blink_strength': deque(maxlen=window_size)
        }
        
        # 统计信息
        self.total_processed = 0
        self.eog_detected = 0
        self.artifacts_removed = 0
        
        logger.info(f"改进EOG伪迹去除器初始化完成 - 阈值:{blink_threshold}, 窗口:{window_size}, ICA成分:{ica_components}")
    
    def process_signal(self, brain_wave_result) -> BrainWaveResult:
        """
        处理脑电信号，去除EOG伪迹
        
        Args:
            brain_wave_result: 输入的脑电数据（BrainWaveResult对象或字典）
            
        Returns:
            BrainWaveResult: 处理后的脑电数据
        """
        # 类型验证和自动转换
        if isinstance(brain_wave_result, dict):
            from response.BrainWaveResult import BrainWaveResult
            brain_wave_result = BrainWaveResult.from_dict(brain_wave_result)
            logger.debug("EOGArtifactRemover: 自动将字典转换为BrainWaveResult对象")
        elif not hasattr(brain_wave_result, 'delta'):
            from exception.base_exceptions import SignalProcessingException
            raise SignalProcessingException(
                f"EOGArtifactRemover: 输入数据类型无效: {type(brain_wave_result)}",
                processing_step="eog_input_validation"
            )
        
        self.total_processed += 1
        
        # 更新缓冲区
        self._update_buffer(brain_wave_result)
        
        # 检查是否有足够的数据进行处理
        if len(self.signal_buffer['delta']) < self.window_size:
            return brain_wave_result
        
        # 获取当前眨眼强度
        current_blink_strength = list(self.signal_buffer['blink_strength'])[-1]
        
        # 检查是否需要进行EOG处理
        if current_blink_strength > self.blink_threshold:
            self.eog_detected += 1
            logger.info(f"【改进EOG】检测到眨眼伪迹 (强度: {current_blink_strength:.1f} > {self.blink_threshold})")
            
            # 尝试去除EOG伪迹
            cleaned_result = self._remove_eog_artifacts_improved(brain_wave_result)
            if cleaned_result is not None:
                self.artifacts_removed += 1
                logger.info("【改进EOG】成功去除EOG伪迹")
                return cleaned_result
            else:
                logger.info("【改进EOG】EOG去除失败，返回原始数据")
        
        return brain_wave_result
    
    def _update_buffer(self, brain_wave_result: BrainWaveResult):
        """更新信号缓冲区（8频域支持）"""
        self.signal_buffer['delta'].append(brain_wave_result.delta)
        self.signal_buffer['theta'].append(brain_wave_result.theta)
        
        # 8频域数据
        self.signal_buffer['low_alpha'].append(brain_wave_result.low_alpha)
        self.signal_buffer['high_alpha'].append(brain_wave_result.high_alpha)
        self.signal_buffer['low_beta'].append(brain_wave_result.low_beta)
        self.signal_buffer['high_beta'].append(brain_wave_result.high_beta)
        self.signal_buffer['low_gamma'].append(brain_wave_result.low_gamma)
        self.signal_buffer['mid_gamma'].append(brain_wave_result.mid_gamma)
        
        # 向后兼容数据
        self.signal_buffer['alpha'].append(brain_wave_result.alpha)
        self.signal_buffer['beta'].append(brain_wave_result.beta)
        self.signal_buffer['gamma'].append(brain_wave_result.gamma)
        
        # 计算眨眼强度
        blink_strength = 0
        if 'processing_info' in brain_wave_result.__dict__ and 'blink_strength' in brain_wave_result.processing_info:
            blink_strength = brain_wave_result.processing_info.get('blink_strength', 0)
        else:
            # 使用delta波段的变化率作为替代指标
            delta_values = list(self.signal_buffer['delta'])
            if len(delta_values) >= 3:
                delta_diff = abs(delta_values[-1] - np.mean(delta_values[-4:-1] if len(delta_values) >= 4 else delta_values[:-1]))
                blink_strength = delta_diff
        
        self.signal_buffer['blink_strength'].append(blink_strength)
    
    def _remove_eog_artifacts_improved(self, brain_wave_result: BrainWaveResult) -> Optional[BrainWaveResult]:
        """
        改进的EOG伪迹去除方法，增强数值稳定性
        
        Args:
            brain_wave_result: 当前脑电数据
            
        Returns:
            BrainWaveResult: 去除EOG后的数据，失败时返回None
        """
        try:
            # 构建8频域多通道信号矩阵
            signal_matrix = np.array([
                list(self.signal_buffer['delta']),
                list(self.signal_buffer['theta']),
                list(self.signal_buffer['low_alpha']),
                list(self.signal_buffer['high_alpha']),
                list(self.signal_buffer['low_beta']),
                list(self.signal_buffer['high_beta']),
                list(self.signal_buffer['low_gamma']),
                list(self.signal_buffer['mid_gamma'])
            ], dtype=np.float64)  # 使用双精度
            
            # 检查数据有效性
            if signal_matrix.shape[0] < 2 or signal_matrix.shape[1] < 10:
                logger.warning("信号矩阵维度不足，无法进行ICA分析")
                return None
            
            # 强化数据预处理
            processed_matrix, original_mean, original_std = self._robust_preprocessing(signal_matrix)
            
            if processed_matrix is None:
                logger.warning("数据预处理失败")
                return None
            
            # 应用改进的ICA分离
            logger.info("【改进EOG】开始ICA分离...")
            ica_sources = self.ica_model.fit_transform(processed_matrix.T).T
            mixing_matrix = self.ica_model.mixing_
            
            # 识别EOG成分
            eog_component_idx = self._identify_eog_component_improved(ica_sources)
            
            if eog_component_idx is not None:
                # 去除EOG成分
                logger.info(f"【改进EOG】去除识别到的EOG成分 {eog_component_idx}")
                cleaned_sources = ica_sources.copy()
                cleaned_sources[eog_component_idx, :] = 0
                
                # 重构信号
                cleaned_signals = (mixing_matrix @ cleaned_sources).T # Transpose back
                
                # 反标准化
                cleaned_signals_denorm = cleaned_signals * original_std.T + original_mean.T

                # 创建清理后的结果
                cleaned_result = brain_wave_result.copy()
                
                # 更新8频域数据（窗口的最后一个值）
                cleaned_result.delta = max(0, int(cleaned_signals_denorm[-1, 0]))
                cleaned_result.theta = max(0, int(cleaned_signals_denorm[-1, 1]))
                cleaned_result.low_alpha = max(0, int(cleaned_signals_denorm[-1, 2]))
                cleaned_result.high_alpha = max(0, int(cleaned_signals_denorm[-1, 3]))
                cleaned_result.low_beta = max(0, int(cleaned_signals_denorm[-1, 4]))
                cleaned_result.high_beta = max(0, int(cleaned_signals_denorm[-1, 5]))
                cleaned_result.low_gamma = max(0, int(cleaned_signals_denorm[-1, 6]))
                cleaned_result.mid_gamma = max(0, int(cleaned_signals_denorm[-1, 7]))
                
                # 重新计算向后兼容的合成频段
                cleaned_result.alpha = cleaned_result.low_alpha + cleaned_result.high_alpha
                cleaned_result.beta = cleaned_result.low_beta + cleaned_result.high_beta
                cleaned_result.gamma = cleaned_result.low_gamma + cleaned_result.mid_gamma
                
                # 重新计算分析指标
                cleaned_result.update_analysis_indicators()
                
                # 输出变化情况
                logger.info(f"【改进EOG】信号变化: Delta: {brain_wave_result.delta} -> {cleaned_result.delta}")
                
                return cleaned_result
            else:
                logger.info("【改进EOG】未识别到明显的EOG成分，跳过处理")
                return None
                
        except Exception as e:
            logger.error(f"【改进EOG】EOG伪迹去除失败: {str(e)}")
            return None
    
    def _robust_preprocessing(self, signal_matrix: np.ndarray) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]:
        """
        强化的数据预处理，提高数值稳定性
        
        Args:
            signal_matrix: 输入信号矩阵
            
        Returns:
            预处理后的信号矩阵，原始均值，原始标准差；失败时返回 (None, None, None)
        """
        try:
            # 检查并处理异常值
            if np.any(np.isinf(signal_matrix)) or np.any(np.isnan(signal_matrix)):
                logger.warning("检测到无穷值或NaN，进行修复")
                signal_matrix = np.nan_to_num(signal_matrix, nan=0.0, posinf=1e6, neginf=-1e6)
            
            # 记录原始均值和标准差
            original_mean = np.mean(signal_matrix, axis=1, keepdims=True)
            original_std = np.std(signal_matrix, axis=1, keepdims=True)
            original_std[original_std < 1e-10] = 1.0

            # 去除直流分量
            signal_matrix_centered = signal_matrix - original_mean
            
            # 标准化
            signal_matrix_normalized = signal_matrix_centered / original_std
            
            # 限制数值范围，避免极值
            signal_matrix_normalized = np.clip(signal_matrix_normalized, -10, 10)
            
            # 应用低通滤波去除高频噪声
            for i in range(signal_matrix_normalized.shape[0]):
                try:
                    # 设计低通滤波器
                    nyquist = self.sampling_rate / 2
                    cutoff = min(100, nyquist * 0.8)  # 截止频率
                    b, a = signal.butter(4, cutoff / nyquist, btype='low')
                    signal_matrix_normalized[i, :] = signal.filtfilt(b, a, signal_matrix_normalized[i, :])
                except Exception as filter_e:
                    logger.warning(f"滤波失败，跳过通道 {i}: {str(filter_e)}")
            
            # 最终检查
            if np.any(np.isinf(signal_matrix_normalized)) or np.any(np.isnan(signal_matrix_normalized)):
                logger.error("预处理后仍存在无效值")
                return None, None, None
            
            return signal_matrix_normalized, original_mean, original_std
            
        except Exception as e:
            logger.error(f"数据预处理失败: {str(e)}")
            return None, None, None
    
    def _identify_eog_component_improved(self, ica_sources: np.ndarray) -> Optional[int]:
        """
        改进的EOG成分识别方法
        
        Args:
            ica_sources: ICA分离的源信号
            
        Returns:
            EOG成分的索引，未找到时返回None
        """
        try:
            eog_scores = []
            
            for i in range(ica_sources.shape[0]):
                component = ica_sources[i, :]
                
                # 计算多个特征
                # 1. 低频能量比例（EOG主要在低频）
                fft_component = np.fft.fft(component)
                freqs = np.fft.fftfreq(len(component), 1/self.sampling_rate)
                low_freq_power = np.sum(np.abs(fft_component[np.abs(freqs) <= 4]))
                total_power = np.sum(np.abs(fft_component))
                low_freq_ratio = low_freq_power / (total_power + 1e-10)
                
                # 2. 峰值特征（眨眼产生的尖峰）
                peak_ratio = np.max(np.abs(component)) / (np.std(component) + 1e-10)
                
                # 3. 偏度（眨眼信号通常有偏度）
                skewness = abs(np.mean(((component - np.mean(component)) / (np.std(component) + 1e-10))**3))
                
                # 综合评分
                eog_score = low_freq_ratio * 0.4 + min(peak_ratio/10, 1.0) * 0.4 + min(skewness, 1.0) * 0.2
                eog_scores.append(eog_score)
                
                logger.debug(f"成分 {i}: 低频比例={low_freq_ratio:.3f}, 峰值比例={peak_ratio:.3f}, 偏度={skewness:.3f}, 总分={eog_score:.3f}")
            
            # 选择得分最高的成分
            best_component = np.argmax(eog_scores)
            best_score = eog_scores[best_component]
            
            # 降低阈值，提高检测敏感性
            if best_score > 0.2:  # 从0.3降低到0.2
                logger.info(f"【改进EOG】识别到EOG成分 {best_component}，得分: {best_score:.3f}")
                return best_component
            else:
                logger.info(f"【改进EOG】未找到明显EOG成分，最高得分: {best_score:.3f}")
                return None
                
        except Exception as e:
            logger.error(f"EOG成分识别失败: {str(e)}")
            return None
    
    def get_statistics(self) -> Dict:
        """获取处理统计信息"""
        detection_rate = (self.eog_detected / self.total_processed * 100) if self.total_processed > 0 else 0
        removal_rate = (self.artifacts_removed / self.eog_detected * 100) if self.eog_detected > 0 else 0
        
        return {
            'total_processed': self.total_processed,
            'eog_detected': self.eog_detected,
            'artifacts_removed': self.artifacts_removed,
            'detection_rate': detection_rate,
            'removal_rate': removal_rate,
            'current_threshold': self.blink_threshold
        }
    
    def get_remover_info(self) -> dict:
        """
        获取伪迹去除器配置信息
        """
        return {
            'blink_threshold': self.blink_threshold,
            'window_size': self.window_size,
            'ica_components': self.ica_components,
            'sampling_rate': self.sampling_rate
        }

    def process_8bands(self, brain_wave_result) -> dict:
        """
        处理8频域脑波数据的EOG伪迹去除
        
        Args:
            brain_wave_result: 包含8频域数据的BrainWaveResult对象或字典
            
        Returns:
            dict: 经过EOG伪迹去除处理的结果字典
        """
        try:
            # 类型验证和转换
            if hasattr(brain_wave_result, 'to_dict'):
                # BrainWaveResult对象，转换为字典
                input_data = brain_wave_result.to_dict()
            elif isinstance(brain_wave_result, dict):
                # 已经是字典
                input_data = brain_wave_result.copy()
            else:
                logger.error(f"EOGArtifactRemover.process_8bands: 不支持的输入类型: {type(brain_wave_result)}")
                return {}
            
            # 创建结果字典
            processed_result = input_data.copy()
            
            # 收集8频域数据
            band_names = ['delta', 'theta', 'low_alpha', 'high_alpha', 
                         'low_beta', 'high_beta', 'low_gamma', 'mid_gamma']
            
            band_data = np.array([
                input_data.get('delta', 0),
                input_data.get('theta', 0),
                input_data.get('low_alpha', 0),
                input_data.get('high_alpha', 0),
                input_data.get('low_beta', 0),
                input_data.get('high_beta', 0),
                input_data.get('low_gamma', 0),
                input_data.get('mid_gamma', 0)
            ])

            # 伪迹检测与去除
            if self.is_blink_detected_8bands(band_data):
                # 使用ICA进行伪迹去除
                corrected_data = self.remove_eog_ica_8bands(band_data)
                
                # 更新字典数据
                if corrected_data is not None:
                    for i, band_name in enumerate(band_names):
                        processed_result[band_name] = int(corrected_data[i])

                    # 添加处理记录
                    if 'processing_info' not in processed_result:
                        processed_result['processing_info'] = {}
                    processed_result['processing_info']['eog_artifact_remover_8bands'] = {
                        'status': 'corrected',
                        'details': 'EOG artifact detected and corrected using ICA.'
                    }
                else:
                    if 'processing_info' not in processed_result:
                        processed_result['processing_info'] = {}
                    processed_result['processing_info']['eog_artifact_remover_8bands'] = {
                        'status': 'correction_failed',
                        'details': 'EOG artifact detected but correction failed.'
                    }
            else:
                if 'processing_info' not in processed_result:
                    processed_result['processing_info'] = {}
                processed_result['processing_info']['eog_artifact_remover_8bands'] = {
                    'status': 'no_artifact',
                    'details': 'No significant EOG artifact detected.'
                }
            
            return processed_result

        except Exception as e:
            logger.error(f"8频域EOG伪迹去除处理失败: {e}")
            return brain_wave_result

    def process_8bands(self, brain_wave_result: BrainWaveResult) -> BrainWaveResult:
        """
        处理8频域脑波数据的EOG伪迹去除
        
        Args:
            brain_wave_result: 包含8频域数据的BrainWaveResult对象
            
        Returns:
            BrainWaveResult: 经过EOG伪迹去除处理的结果
        """
        try:
            # 创建结果副本
            processed_result = brain_wave_result.copy()
            
            # 收集8频域数据
            band_data = np.array([
                processed_result.delta,
                processed_result.theta,
                processed_result.low_alpha,
                processed_result.high_alpha,
                processed_result.low_beta,
                processed_result.high_beta,
                processed_result.low_gamma,
                processed_result.mid_gamma
            ])

            # 伪迹检测与去除
            if self.is_blink_detected_8bands(band_data):
                # 使用ICA进行伪迹去除
                corrected_data = self.remove_eog_ica_8bands(band_data)
                
                # 更新BrainWaveResult
                if corrected_data is not None:
                    processed_result.delta = int(corrected_data[0])
                    processed_result.theta = int(corrected_data[1])
                    processed_result.low_alpha = int(corrected_data[2])
                    processed_result.high_alpha = int(corrected_data[3])
                    processed_result.low_beta = int(corrected_data[4])
                    processed_result.high_beta = int(corrected_data[5])
                    processed_result.low_gamma = int(corrected_data[6])
                    processed_result.mid_gamma = int(corrected_data[7])

                    # 添加处理记录
                    processed_result.add_processing_step('eog_artifact_remover_8bands', {
                        'status': 'corrected',
                        'details': 'EOG artifact detected and corrected using ICA.'
                    })
                else:
                    processed_result.add_processing_step('eog_artifact_remover_8bands', {
                        'status': 'correction_failed',
                        'details': 'EOG artifact detected but correction failed.'
                    })
            else:
                processed_result.add_processing_step('eog_artifact_remover_8bands', {
                    'status': 'no_artifact',
                    'details': 'No significant EOG artifact detected.'
                })

            # 更新分析指标
            processed_result.update_analysis_indicators()
            
            return processed_result

        except Exception as e:
            logger.error(f"8频域EOG伪迹去除处理失败: {e}")
            return brain_wave_result
    
    def reset_statistics(self):
        """重置统计信息"""
        self.total_processed = 0
        self.eog_detected = 0
        self.artifacts_removed = 0
        
        # 清空缓冲区
        for buffer in self.signal_buffer.values():
            buffer.clear()